import numpy as np

COCO_PERSON_SKELETON = [(15, 14), (14, 13), (12, 11), (11, 10), (13, 10), (6, 13), (3, 10), (6, 3), (6, 7), (3, 4),
                        (7, 8), (4, 5), (17, 16), (1, 17), (1, 16), (17, 19), (16, 18), (19, 6), (18, 3), (1, 2),
                        (2, 3), (2, 6), (2, 9), (9, 10), (12, 25), (12, 23), (23, 24), (9, 13), (15, 22), (15, 20),
                        (20, 21)]

KINEMATIC_TREE_SKELETON = [(2, 9), (2, 3), (2, 6), (3, 4), (4, 5), (6, 7), (7, 8), (9, 10), (10, 11), (11, 12), (9, 13),
                           (13, 14), (14, 15), (2, 1), (1, 16), (16, 18), (1, 17), (17, 19), (3, 18), (6, 19), (15, 20),
                           (20, 21), (15, 22), (12, 23), (23, 24), (12, 25)]
COCO_KEYPOINTS = [
    "nose",  # 1
    "Neck",  # 2
    "right_shoulder",  # 3
    "right_elbow",  # 4
    "right_wrist",  # 5
    "left_shoulder",  # 6
    "left_elbow",  # 7
    "left_wrist",  # 8
    "MidHip",  # 9
    "right_hip",  # 10
    "right_knee",  # 11
    "right_ankle",  # 12
    "left_hip",  # 13
    "left_knee",  # 14
    "left_ankle",  # 15
    "right_eye",  # 16
    "left_eye",  # 17
    "right_ear",  # 18
    "left_ear",  # 19
    "LBigToe",  # 20
    "LSmallToe",  # 21
    "LHeel",  # 22
    "RBigToe",  # 23
    "RSmallToe",  # 24
    "RHeel"  # 25
]

COCO_UPRIGHT_POSE = np.array(
    [[0.0, 10.3, 2.0], [0.0, 9.0, 2.0], [1.4, 9.0, 2.0], [1.75, 7.2, 2.0], [1.75, 5.2, 2.0], [-1.4, 9.0, 2.0],
     [-1.75, 7.0, 2.0], [-1.75, 5.0, 2.0], [0.0, 5.0, 2.0], [1.26, 5.0, 2.0], [1.4, 3.1, 2.0], [1.4, 1.0, 2.0],
     [-1.26, 5.0, 2.0], [-1.4, 3.0, 2.0], [-1.4, 1.0, 2.0], [0.35, 10.7, 2.0], [-0.35, 10.7, 2.0], [0.7, 10.5, 2.0],
     [-0.7, 10.5, 2.0], [-1.75, 0.0, 2.0], [-2.0, 0.25, 2.0], [-1.15, 0.75, 2.0], [1.75, 0.0, 2.0], [2.0, 0.25, 2.0],
     [1.15, 0.75, 2.0]])

COCO_DAVINCI_POSE = np.array(
    [[0.0, 9.3, 2.0], [-1.26, 4.0, 2.0], [1.4, 8.0, 2.0], [3.3, 9.2, 2.0], [4.5, 10.7, 2.0], [-1.4, 8.0, 2.0],
     [-3.3, 9.0, 2.0], [-4.5, 10.5, 2.0], [-1.26, 4.0, 2.0], [1.26, 4.0, 2.0], [2.0, 2.1, 2.0], [2.4, 0.1, 2.0],
     [-1.26, 4.0, 2.0], [-2.0, 2.0, 2.0], [-2.4, 0.0, 2.0], [0.35, 9.7, 2.0], [-0.35, 9.7, 2.0], [0.7, 9.5, 2.0],
     [-0.7, 9.5, 2.0], [-1.26, 4.0, 2.0], [-1.26, 4.0, 2.0], [-1.26, 4.0, 2.0], [-1.26, 4.0, 2.0], [-1.26, 4.0, 2.0],
     [-1.26, 4.0, 2.0]])

HFLIP = {
    'left_eye': 'right_eye',
    'right_eye': 'left_eye',
    'left_ear': 'right_ear',
    'right_ear': 'left_ear',
    'left_shoulder': 'right_shoulder',
    'right_shoulder': 'left_shoulder',
    'left_elbow': 'right_elbow',
    'right_elbow': 'left_elbow',
    'left_wrist': 'right_wrist',
    'right_wrist': 'left_wrist',
    'left_hip': 'right_hip',
    'right_hip': 'left_hip',
    'left_knee': 'right_knee',
    'right_knee': 'left_knee',
    'left_ankle': 'right_ankle',
    'right_ankle': 'left_ankle',
    'LBigToe': 'RBigToe',
    'RBigToe': 'LBigToe',
    'LSmallToe': 'RSmallToe',
    'RSmallToe': 'LSmallToe',
    'LHeel': 'RHeel',
    'RHeel': 'LHeel',
}
DENSER_COCO_PERSON_SKELETON = [(1, 17), (1, 16), (17, 16), (1, 19), (1, 18), (19, 18), (1, 6), (1, 3), (17, 6), (16, 3),
                               (17, 19), (16, 18), (19, 6), (18, 3), (6, 3), (6, 13), (3, 10), (6, 10), (3, 13),
                               (13, 10), (6, 7), (3, 4), (7, 8), (4, 5), (6, 8), (3, 5), (7, 4), (8, 5), (8, 13),
                               (5, 10), (8, 14), (5, 11), (14, 13), (11, 10), (13, 11), (10, 14), (13, 15), (10, 12),
                               (15, 14), (12, 11), (14, 12), (11, 15), (14, 11), (15, 12), (1, 2), (2, 3), (2, 6),
                               (2, 9), (9, 10), (12, 25), (12, 23), (23, 24), (9, 13), (15, 22), (15, 20), (20, 21),
                               (12, 15), (25, 22), (23, 20), (24, 21)]

DENSER_COCO_PERSON_CONNECTIONS = [
    c
    for c in DENSER_COCO_PERSON_SKELETON
    if c not in COCO_PERSON_SKELETON
]

COCO_PERSON_SIGMAS = [0.026, 0.025, 0.079, 0.072, 0.062, 0.079, 0.072, 0.062, 0.025, 0.107, 0.087, 0.089, 0.107, 0.087,
                      0.089, 0.025, 0.025, 0.035, 0.035, 0.025, 0.025, 0.025, 0.025, 0.025, 0.025]
COCO_CATEGORIES = [
    'person',
    'bicycle',
    'car',
    'motorcycle',
    'airplane',
    'bus',
    'train',
    'truck',
    'boat',
    'traffic light',
    'fire hydrant',
    'street sign',
    'stop sign',
    'parking meter',
    'bench',
    'bird',
    'cat',
    'dog',
    'horse',
    'sheep',
    'cow',
    'elephant',
    'bear',
    'zebra',
    'giraffe',
    'hat',
    'backpack',
    'umbrella',
    'shoe',
    'eye glasses',
    'handbag',
    'tie',
    'suitcase',
    'frisbee',
    'skis',
    'snowboard',
    'sports ball',
    'kite',
    'baseball bat',
    'baseball glove',
    'skateboard',
    'surfboard',
    'tennis racket',
    'bottle',
    'plate',
    'wine glass',
    'cup',
    'fork',
    'knife',
    'spoon',
    'bowl',
    'banana',
    'apple',
    'sandwich',
    'orange',
    'broccoli',
    'carrot',
    'hot dog',
    'pizza',
    'donut',
    'cake',
    'chair',
    'couch',
    'potted plant',
    'bed',
    'mirror',
    'dining table',
    'window',
    'desk',
    'toilet',
    'door',
    'tv',
    'laptop',
    'mouse',
    'remote',
    'keyboard',
    'cell phone',
    'microwave',
    'oven',
    'toaster',
    'sink',
    'refrigerator',
    'blender',
    'book',
    'clock',
    'vase',
    'scissors',
    'teddy bear',
    'hair drier',
    'toothbrush',
    'hair brush',
]


def draw_ann(ann, *, keypoint_painter, filename=None, margin=0.5, aspect=None, **kwargs):
    from .. import show  # pylint: disable=import-outside-toplevel

    bbox = ann.bbox()
    xlim = bbox[0] - margin, bbox[0] + bbox[2] + margin
    ylim = bbox[1] - margin, bbox[1] + bbox[3] + margin
    if aspect == 'equal':
        fig_w = 5.0
    else:
        fig_w = 5.0 / (ylim[1] - ylim[0]) * (xlim[1] - xlim[0])

    with show.canvas(filename, figsize=(fig_w, 5), nomargin=True, **kwargs) as ax:
        ax.set_axis_off()
        ax.set_xlim(*xlim)
        ax.set_ylim(*ylim)

        if aspect is not None:
            ax.set_aspect(aspect)

        keypoint_painter.annotation(ax, ann)


def draw_skeletons(pose):
    from ..annotation import Annotation  # pylint: disable=import-outside-toplevel
    from .. import show  # pylint: disable=import-outside-toplevel

    scale = np.sqrt(
        (np.max(pose[:, 0]) - np.min(pose[:, 0]))
        * (np.max(pose[:, 1]) - np.min(pose[:, 1]))
    )

    show.KeypointPainter.show_joint_scales = True
    keypoint_painter = show.KeypointPainter(color_connections=True, linewidth=6)

    ann = Annotation(keypoints=COCO_KEYPOINTS, skeleton=COCO_PERSON_SKELETON)
    ann.set(pose, np.array(COCO_PERSON_SIGMAS) * scale)
    draw_ann(ann, filename='docs/skeleton_coco.png', keypoint_painter=keypoint_painter)

    ann = Annotation(keypoints=COCO_KEYPOINTS, skeleton=KINEMATIC_TREE_SKELETON)
    ann.set(pose, np.array(COCO_PERSON_SIGMAS) * scale)
    draw_ann(ann, filename='docs/skeleton_kinematic_tree.png', keypoint_painter=keypoint_painter)

    ann = Annotation(keypoints=COCO_KEYPOINTS, skeleton=DENSER_COCO_PERSON_SKELETON)
    ann.set(pose, np.array(COCO_PERSON_SIGMAS) * scale)
    draw_ann(ann, filename='docs/skeleton_dense.png', keypoint_painter=keypoint_painter)


def print_associations():
    for j1, j2 in COCO_PERSON_SKELETON:
        print(COCO_KEYPOINTS[j1 - 1], '-', COCO_KEYPOINTS[j2 - 1])


if __name__ == '__main__':
    print_associations()

    # c, s = np.cos(np.radians(45)), np.sin(np.radians(45))
    # rotate = np.array(((c, -s), (s, c)))
    # rotated_pose = np.copy(COCO_DAVINCI_POSE)
    # rotated_pose[:, :2] = np.einsum('ij,kj->ki', rotate, rotated_pose[:, :2])
    draw_skeletons(COCO_UPRIGHT_POSE)
